set more off

* range of years in analysis 
local yearMin = 1981
local initial_date = "01/01/`yearMin'"
local yearMax = 2012

* Parameters for water-balance model
local T_rain = 3.3
local T_snow = -10
local drofrac = 0.05
local meltmax = 0.5

* Define soil moisture and snow storage in initial period
use "..\dataRaw\county_soils_usa", clear
keep stcofips rootznaws
gen date=date("`initial_date'", "MDY") - 1
********************************
gen soil_moisture_usgs=0.5*rootznaws
gen soil_moisture=0.5*rootznaws
gen soil_moisture_corn=0.5*rootznaws
gen snostor=0
********************************
format date %td
save "..\temp\initial_soil_moisture", replace

* Import file of crop coefficients
import excel using "..\dataRaw\crop coefficients ET.xlsx", sheet("data") first clear
save "..\dataRaw\crop coefficients ET", replace


forvalues t = `yearMin'/`yearMax' {
		display "`t'"
qui{		
        ******************************************************************************
        * load data
        use "..\dataRaw\PRISM_county_`t'", clear 
		merge 1:1 stcofips year month day using "..\dataRaw\PRISM_county_moreDday_`t'", nogen
		
		gen date_str=string(month) + "/" + string(day) + "/" + string(year)
		gen date=date(date_str,"MDY")
		drop date_str
		format date %td
		xtset stcofips date

		merge m:1 stcofips using "..\dataRaw\county_soils_usa", keep(match master) keepusing(rootznaws) nogen
		
		* Snow Accumulation
		gen ppt_snow=ppt if tavg<=`T_snow'
		replace ppt_snow=ppt*(`T_rain' - tavg)/(`T_rain'-`T_snow') if tavg>`T_snow' & tavg<`T_rain'
		replace ppt_snow=0 if tavg>=`T_rain'
		gen ppt_rain=ppt-ppt_snow

		* Direct Runoff
		gen ppt_remain=ppt_rain - ppt_rain*`drofrac'

		* Snow Melt
		gen smf=0 if tavg<=`T_snow'
		replace smf=(tavg - `T_snow')/(`T_rain'-`T_snow')*`meltmax' if tavg>`T_snow' & tavg<`T_rain'
		replace smf=`meltmax' if tavg>=`T_rain'
		
		* Generate corn-specific ET
		gen plant_date=mdy(5,1,year)
		gen Days_after_plant=date-plant_date
		merge m:1 Days_after_plant using "..\dataRaw\crop coefficients ET", nogen keep(master match)
		replace kc=0.25 if kc==.
		replace kc=0.25 if kc<0.25
		gen ETc=ET0*kc
		drop plant_date Days_after_plant
		sort stcofips date
	
		gen aet=.
		gen soil_moisture=.
		gen water_deficit=.
		gen water_surplus=.
		
		gen snostor=.
		gen sm=.
		gen ppt_total=.
		gen stw=.
		gen aet_usgs=.
		gen soil_moisture_usgs=.
		gen water_deficit_usgs=.
		gen water_surplus_usgs=.	
		gen aet_corn=.
		gen soil_moisture_corn=.
		gen water_deficit_corn=.
		gen water_surplus_corn=.
		
		* append initial soil moisture and snow storage
		if `t'==`yearMin' {
			append using "..\temp\initial_soil_moisture"
		}
		else {
			local tMinus1=`t'-1
			append using `WaterBalanceDec31_`tMinus1''
		}
		
		sort stcofips date
		replace snostor=(ppt_snow + l.snostor)*(1-smf) if date>=date("01/01/`t'", "MDY")
		replace sm=(ppt_snow + l.snostor)*smf if date>=date("01/01/`t'", "MDY")
		replace ppt_total=ppt_remain + sm if date>=date("01/01/`t'", "MDY")
		
		local start = date("01/01/`t'", "MDY")
		local end = date("12/31/`t'", "MDY")
		forvalues day = `start'/`end'  {
			***********************
			* Simple Water Balance
			***********************
			replace soil_moisture=l.soil_moisture + ppt - ET0 if date==`day'
			replace soil_moisture=0 if soil_moisture<0 & date==`day'
			replace water_surplus=soil_moisture - rootznaws if soil_moisture>rootznaws & date==`day'
			replace soil_moisture=rootznaws if soil_moisture>rootznaws & date==`day'
			* If precip < ET0
			replace aet=ppt - d.soil_moisture if ppt<ET0 & date==`day'
			* If precip >=ET0
			replace aet=ET0 if ppt>=ET0 & date==`day'
			replace water_deficit=ET0-aet if aet<ET0 & date==`day'
			
			***********************
			* USGS Water Balance
			***********************
			* If precip < ET0
			replace stw=(ET0-ppt_total)*l.soil_moisture_usgs/rootznaws if ppt_total<ET0 & date==`day'
			replace aet_usgs=ppt_total + stw if ppt_total<ET0 & date==`day'
			replace soil_moisture_usgs=l.soil_moisture_usgs - stw if ppt_total<ET0 & date==`day'
			replace soil_moisture_usgs=0 if soil_moisture_usgs<0 & date==`day'
			replace water_deficit_usgs=ET0-aet_usgs if aet_usgs<ET0 & date==`day'
	
			* If precip >= ET0
			replace aet_usgs=ET0 if ppt_total>=ET0 & date==`day'
			replace soil_moisture_usgs=l.soil_moisture_usgs + ppt_total-ET0 if ppt_total>=ET0 & date==`day'
			replace water_surplus_usgs=soil_moisture_usgs - rootznaws if soil_moisture_usgs>rootznaws & date==`day'
			replace soil_moisture_usgs=rootznaws if soil_moisture_usgs>rootznaws & date==`day'
			
			***********************
			* Corn Water Balance
			***********************
			replace soil_moisture_corn=l.soil_moisture_corn + ppt - ETc if date==`day'
			replace soil_moisture_corn=0 if soil_moisture_corn<0 & date==`day'
			replace water_surplus_corn=soil_moisture_corn - rootznaws if soil_moisture_corn>rootznaws & date==`day'
			replace soil_moisture_corn=rootznaws if soil_moisture_corn>rootznaws & date==`day'
			* If precip < ETc
			replace aet_corn=ppt - d.soil_moisture_corn if ppt<ETc & date==`day'
			* If precip >=ETc
			replace aet_corn=ETc if ppt>=ETc & date==`day'
			replace water_deficit_corn=ETc-aet_corn if aet_corn<ETc & date==`day'
			
		} //end of looping over days
		replace water_deficit=0 if water_deficit==. & rootznaws!=. & ET0!=.
		replace water_surplus=0 if water_surplus==. & rootznaws!=. & ET0!=.
		replace water_deficit_usgs=0 if water_deficit_usgs==. & rootznaws!=. & ET0!=.
		replace water_surplus_usgs=0 if water_surplus_usgs==. & rootznaws!=. & ET0!=.
		replace water_deficit_corn=0 if water_deficit_corn==. & rootznaws!=. & ETc!=.
		replace water_surplus_corn=0 if water_surplus_corn==. & rootznaws!=. & ETc!=.
		
		save "..\temp\PRISM_WaterBalance_`t'", replace
		
		*keep 12/31 for initial period in next year
		keep if date==date("12/31/`t'", "MDY")
		tempfile WaterBalanceDec31_`t'
		save `WaterBalanceDec31_`t''
				
} //end of qui
} //end of looping over years



